#include "trajectory_optimization/trajectory_generator.h"

namespace Planner
{
    using std::cout;
    using std::endl;

    void TrajectoryGenerator::setParams(int order, double speed, int continus)
    {
        order_ = order;
        continus_ = continus;
        speed_ = speed;
        useDefaultTimeSeq = true;
    }

    void TrajectoryGenerator::setWayPoints(const vector<Vector3d> &way_points)
    {
        way_points_ = way_points;
        segments_ = way_points_.size() - 1;
    }

    void TrajectoryGenerator::setBoundaryCondition(const Vector3d &startVel, const Vector3d &startAcc,
                                                   const Vector3d &goalVel, const Vector3d &goalAcc)
    {
        startVel_ = startVel;
        startAcc_ = startAcc;
        goalVel_ = goalVel;
        goalAcc_ = goalAcc;
    }

    void TrajectoryGenerator::numFixFree(int &num_fix, int &num_free)
    {
        num_fix = segments_ + 5;
        num_free = 2 * segments_ - 2;
    }

    void TrajectoryGenerator::solve(int type)
    {
        type_ = type;
        // if (useDefaultTimeSeq)
        //     timeAllocation();
        // useDefaultTimeSeq = true;
        timeAllocation();
        computeQ();
        computeA();
        computeMC();

        MatrixXd invA = inv(A_);
        AiMC_ = invA * MC_;
        R_ = AiMC_.transpose() * Q_ * AiMC_;
        // int num_fix = segments_ + 5;
        // int num_free = 2 * segments_ - 2;
        int num_fix, num_free;
        numFixFree(num_fix, num_free);
        Rpp_ = R_.block(num_fix, num_fix, num_free, num_free);
        Rfp_ = R_.block(0, num_fix, num_fix, num_free);

        // 设置约束变量
        computeFixVal(dfx_, 0);
        computeFixVal(dfy_, 1);
        computeFixVal(dfz_, 2);

        // 闭式求解
        computeSingleAxis(dfx_, coeffsx_, dpx_);
        computeSingleAxis(dfy_, coeffsy_, dpy_);
        computeSingleAxis(dfz_, coeffsz_, dpz_);

        // 计算误差项
        MatrixXd costx = coeffsx_.transpose() * Q_ * coeffsx_;
        MatrixXd costy = coeffsy_.transpose() * Q_ * coeffsy_;
        MatrixXd costz = coeffsz_.transpose() * Q_ * coeffsz_;
        cost_ = costx(0, 0) + costy(0, 0) + costz(0, 0);
        // cout << "initial cost_smooth: " << cost_ << endl;
        // cout << "initial Q:\n";
        // printMatrix(Q_);
    }

    void TrajectoryGenerator::setTimeSeq(vector<double> &time_seq)
    {
        time_seq_ = time_seq;
        useDefaultTimeSeq = false;
    }

    void TrajectoryGenerator::getTimeSeq(vector<double> &time_seq)
    {
        time_seq = time_seq_;
    }

    void TrajectoryGenerator::setCoeffs(vector<MatrixXd> &coeffs)
    {
        if (coeffs.size() != 3)
            cout << "coeffs size error\n";

        coeffsx_ = coeffs[0];
        coeffsy_ = coeffs[1];
        coeffsz_ = coeffs[2];
    }

    void TrajectoryGenerator::getCoeffs(vector<MatrixXd> &coeffs)
    {
        coeffs.clear();
        coeffs.push_back(resize(coeffsx_));
        coeffs.push_back(resize(coeffsy_));
        coeffs.push_back(resize(coeffsz_));
    }

    void TrajectoryGenerator::timeAllocation()
    {
        time_seq_.clear();
        time_seq_.reserve(segments_);
        for (size_t i = 0; i < segments_; i++)
        {
            double len = (way_points_[i + 1] - way_points_[i]).norm();
            double duration = len / speed_;
            time_seq_.push_back(duration);
        }
    }

    void TrajectoryGenerator::computeQ()
    {
        Q_ = MatrixXd::Zero(segments_ * (order_ + 1), segments_ * (order_ + 1));
        MatrixXd Qi;
        for (size_t i = 0; i < segments_; i++)
        {
            computeQi(Qi, time_seq_[i]);
            Q_.block(i * (order_ + 1), i * (order_ + 1), order_ + 1, order_ + 1) = Qi;
        }
    }

    void TrajectoryGenerator::computeQi(MatrixXd &Qi, double time_seq)
    {
        Qi = MatrixXd::Zero(order_ + 1, order_ + 1);
        for (size_t i = type_; i <= order_; i++)
            for (size_t j = i; j <= order_; j++)
            {
                int k1 = i - type_;
                int k2 = j - type_;
                int k = k1 + k2 + 1;
                Qi(i, j) = prod(k1 + type_, k1 + 1) * prod(k2 + type_, k2 + 1) / double(k) * std::pow(time_seq, k);
                Qi(j, i) = Qi(i, j);
            }
    }

    double TrajectoryGenerator::prod(int m, int n)
    {
        int sum = 1;
        for (size_t i = n; i <= m; i++)
            sum *= i;
        return double(sum);
    }

    void TrajectoryGenerator::computeA()
    {
        A_ = MatrixXd::Zero(2 * continus_ * segments_, segments_ * (order_ + 1));
        MatrixXd Ai;
        // d中排列为p1(t0)\p1(t1)\v1(t0)\v1(t1)\a0(t0)\a0(t1)
        for (size_t i = 0; i < segments_; i++) // 轨迹段数
        {
            Ai = MatrixXd::Zero(2 * continus_, order_ + 1);
            for (size_t j = 0; j < continus_; j++) // j表示p\v\a
            {
                Ai(2 * j, j) = prod(j);
                for (size_t k = j; k < order_ + 1; k++) // 多项式系数
                    Ai(2 * j + 1, k) = prod(k) / prod(k - j) * std::pow(time_seq_[i], k - j);
            }
            A_.block(2 * continus_ * i, i * (order_ + 1), 2 * continus_, order_ + 1) = Ai;
        }
    }

    // 公式为d = MC [dfix dfree] M用于消除重复性变量，C为选择矩阵
    // dfix = [p0 p1 ... pn v0 a0 vn an] dfree = [v1 a1 v2 a2 .. vn-1 an-1]
    void TrajectoryGenerator::computeMC()
    {
        MC_ = MatrixXd::Zero(2 * continus_ * segments_, continus_ * (segments_ + 1));
        int num_fix = segments_ + 5;
        int num_free = 2 * segments_ - 2;

        for (size_t i = 0; i < segments_; i++)
        {
            // 固定p
            MC_(6 * i, i) = 1;
            MC_(6 * i + 1, i + 1) = 1;
            if (i != 0 && i != segments_ - 1)
            {
                // 固定v
                MC_(6 * i + 2, num_fix + 2 * (i - 1)) = 1;
                MC_(6 * i + 3, num_fix + 2 * i) = 1;
                // 固定a
                MC_(6 * i + 4, num_fix + 2 * i - 1) = 1;
                MC_(6 * i + 5, num_fix + 2 * i + 1) = 1;
            }
        }
        // 加入边界条件
        MC_(2, segments_ + 1) = 1;
        MC_(3, num_fix) = 1;
        MC_(4, segments_ + 2) = 1;
        MC_(5, num_fix + 1) = 1;
        MC_(6 * (segments_ - 1) + 2, num_fix + num_free - 2) = 1;
        MC_(6 * (segments_ - 1) + 3, segments_ + 3) = 1;
        MC_(6 * (segments_ - 1) + 4, num_fix + num_free - 1) = 1;
        MC_(6 * (segments_ - 1) + 3, segments_ + 4) = 1;
    }

    void TrajectoryGenerator::computeFixVal(MatrixXd &df, int axis)
    {
        df = MatrixXd(segments_ + 5, 1);
        for (size_t i = 0; i <= segments_; i++)
            df(i, 0) = way_points_[i][axis];

        df(segments_ + 1, 0) = startVel_[axis];
        df(segments_ + 2, 0) = startAcc_[axis];
        df(segments_ + 3, 0) = goalVel_[axis];
        df(segments_ + 4, 0) = goalAcc_[axis];
    }

    void TrajectoryGenerator::computeSingleAxis(const MatrixXd &dfix, MatrixXd &coeffs, MatrixXd &dfree)
    {
        dfree = -Rpp_.inverse() * Rfp_.transpose() * dfix;
        MatrixXd dp(3 * (segments_ + 1), 1);
        dp << dfix, dfree;
        coeffs = AiMC_ * dp;
    }

    MatrixXd TrajectoryGenerator::inv(const MatrixXd &matrix)
    {
        MatrixXd inv_matrix;
        if (matrix.rows() != matrix.cols())
            // inv_matrix = matrix.completeOrthogonalDecomposition().pseudoInverse();
            inv_matrix = pseudoInverse(matrix);
        else
            inv_matrix = matrix.inverse();
        return inv_matrix;
    }

    template <typename _Matrix_Type_>
    _Matrix_Type_ TrajectoryGenerator::pseudoInverse(const _Matrix_Type_ &a, double epsilon)
    {
        Eigen::JacobiSVD<_Matrix_Type_> svd(a, Eigen::ComputeThinU | Eigen::ComputeThinV);
        double tolerance = epsilon * std::max(a.cols(), a.rows()) * svd.singularValues().array().abs()(0);
        return svd.matrixV() * (svd.singularValues().array().abs() > tolerance).select(svd.singularValues().array().inverse(), 0).matrix().asDiagonal() * svd.matrixU().adjoint();
    }

    MatrixXd TrajectoryGenerator::resize(const MatrixXd &coeffs)
    {
        MatrixXd coeffs_resized(segments_, order_ + 1);
        MatrixXd temp;
        for (size_t i = 0; i < segments_; i++)
        {
            temp = coeffs.block(i * (order_ + 1), 0, order_ + 1, 1);
            coeffs_resized.row(i) = temp.transpose();
        }
        return coeffs_resized;
    }

    void TrajectoryGenerator::getVals(MatrixXd &R, MatrixXd &Rfp, MatrixXd &Rpp, vector<MatrixXd> &dp,
                                      vector<MatrixXd> &df, MatrixXd &AiMC)
    {
        R = R_;
        Rfp = Rfp_;
        Rpp = Rpp_;
        dp = {dpx_, dpy_, dpz_};
        df = {dfx_, dfy_, dfz_};
        AiMC = AiMC_;
    }

    double TrajectoryGenerator::getCost()
    {
        return cost_;
    }

    void TrajectoryGenerator::solveByDpAndTime(const vector<MatrixXd> &dfree, const vector<double> &time_seq,
                                               vector<MatrixXd> &coeffs, double &cost)
    {
        time_seq_ = time_seq;
        computeQ();
        computeA();
        MatrixXd invA = inv(A_);
        AiMC_ = invA * MC_;
        MatrixXd dpx(3 * (segments_ + 1), 1), dpy(3 * (segments_ + 1), 1), dpz(3 * (segments_ + 1), 1);
        dpx << dfx_, dfree[0];
        dpy << dfy_, dfree[1];
        dpz << dfz_, dfree[2];

        coeffsx_ = AiMC_ * dpx;
        coeffsy_ = AiMC_ * dpy;
        coeffsz_ = AiMC_ * dpz;
        coeffs = {coeffsx_, coeffsy_, coeffsz_};
        MatrixXd costx = coeffsx_.transpose() * Q_ * coeffsx_;
        MatrixXd costy = coeffsy_.transpose() * Q_ * coeffsy_;
        MatrixXd costz = coeffsz_.transpose() * Q_ * coeffsz_;
        cost = costx(0, 0) + costy(0, 0) + costz(0, 0);
    }

    void TrajectoryGenerator::getPos(double t, Vector3d &pos, int seg, bool setSeg)
    {
        pos.setZero();
        if (!setSeg)
        {
            // 如果没有指定轨迹的段数，认为t为总时间戳
            for (size_t i = 0; i < segments_; i++)
            {
                if (t < time_seq_[i])
                    break;
                t -= time_seq_[i];
                seg++;
            }
        }

        // 如果超过了最后的时刻
        if (seg == segments_)
        {
            seg--;
            t = time_seq_.back();
        }

        for (int i = 0; i <= order_; i++)
        {
            pos[0] += coeffsx_(seg * (order_ + 1) + i, 0) * std::pow(t, i);
            pos[1] += coeffsy_(seg * (order_ + 1) + i, 0) * std::pow(t, i);
            pos[2] += coeffsz_(seg * (order_ + 1) + i, 0) * std::pow(t, i);
        }
    }

    void TrajectoryGenerator::getVel(double t, Vector3d &vel, int seg, bool setSeg)
    {
        vel.setZero();
        if (!setSeg)
        {
            // 如果没有指定轨迹的段数，认为t为总时间戳
            for (size_t i = 0; i < segments_; i++)
            {
                if (t < time_seq_[i])
                    break;
                t -= time_seq_[i];
                seg++;
            }
        }

        // 如果超过了最后的时刻
        if (seg == segments_)
        {
            seg--;
            t = time_seq_.back();
        }

        for (int i = 1; i <= order_; i++)
        {
            vel[0] += i * coeffsx_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 1);
            vel[1] += i * coeffsy_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 1);
            vel[2] += i * coeffsz_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 1);
        }
    }

    void TrajectoryGenerator::getAcc(double t, Vector3d &acc, int seg, bool setSeg)
    {
        acc.setZero();
        if (!setSeg)
        {
            // 如果没有指定轨迹的段数，认为t为总时间戳
            for (size_t i = 0; i < segments_; i++)
            {
                if (t < time_seq_[i])
                    break;
                t -= time_seq_[i];
                seg++;
            }
        }

        // 如果超过了最后的时刻
        if (seg == segments_)
        {
            seg--;
            t = time_seq_.back();
        }

        for (int i = 2; i <= order_; i++)
        {
            acc[0] += i * (i - 1) * coeffsx_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 2);
            acc[1] += i * (i - 1) * coeffsy_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 2);
            acc[2] += i * (i - 1) * coeffsz_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 2);
        }
    }

    void TrajectoryGenerator::getState(double t, State &curr_state, int seg, bool setSeg)
    {
        curr_state.pos.setZero();
        curr_state.vel.setZero();
        curr_state.acc.setZero();
        if (!setSeg)
        {
            // 如果没有指定轨迹的段数，认为t为总时间戳
            for (size_t i = 0; i < segments_; i++)
            {
                if (t < time_seq_[i])
                    break;
                t -= time_seq_[i];
                seg++;
            }
        }

        curr_state.t = t;
        curr_state.seg = seg;
        for (int i = 0; i <= order_; i++)
        {
            curr_state.pos[0] += coeffsx_(seg * (order_ + 1) + i, 0) * std::pow(t, i);
            curr_state.pos[1] += coeffsy_(seg * (order_ + 1) + i, 0) * std::pow(t, i);
            curr_state.pos[2] += coeffsz_(seg * (order_ + 1) + i, 0) * std::pow(t, i);
            curr_state.vel[0] += i * coeffsx_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 1);
            curr_state.vel[1] += i * coeffsy_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 1);
            curr_state.vel[2] += i * coeffsz_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 1);
            curr_state.acc[0] += i * (i - 1) * coeffsx_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 2);
            curr_state.acc[1] += i * (i - 1) * coeffsy_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 2);
            curr_state.acc[2] += i * (i - 1) * coeffsz_(seg * (order_ + 1) + i, 0) * std::pow(t, i - 2);
        }
    }
} // namespace Planner